import heapq
import logging
import random
import sys
import threading
import time
import traceback
from typing import Type, TypeVar
import peewee
import typing

from emi import report
from emi.util import frozendict

_DEBUG = True
logger = logging.getLogger('peewee.pool')

PEEWEE_MODEL_TYPE = TypeVar("PEEWEE_MODEL_TYPE", bound=Type[peewee.Model])  # bound表明该类为BaseModel的子类

PeeweeDatabaseWithStatementSingle = typing.NewType('PeeweeDatabaseWithStatementSingle', typing.ContextManager[peewee.Database])

PEEWEE_MODEL= TypeVar("PEEWEE_MODEL", bound=peewee.Model)  # bound表明该类为BaseModel的子类



class PeeweeModelWithStatement():

    def __init__(self,source,read_only:bool):
        self._model_source:PeeweeModelSource[PEEWEE_MODEL_TYPE] = source
        self._read_only = read_only

    def __enter__(self) -> PEEWEE_MODEL_TYPE:
        if self._model_source is None:
            raise RuntimeError("can enter with only once!! ")
        source:DBSource = self._model_source._source
        self._connect_session = source.use_connect(self._read_only)
        db_source = self._connect_session.__enter__()
        peewee_model = self._model_source._get_state_model()
        peewee_model._meta.database = db_source
        return peewee_model

    def __exit__(self, exc_type, exc_val, exc_tb):
        peewee_model = self._model_source._get_state_model()
        peewee_model._meta.database = None
        self._connect_session.__exit__(exc_type,exc_val,exc_tb)
        self._model_source = None
        self._connect_session = None


class _PeeweeModelState(threading.local):

    def __init__(self, **kwargs):
        super(_PeeweeModelState, self).__init__(**kwargs)
        self.model = None


class PeeweeModelSource(typing.Generic[PEEWEE_MODEL_TYPE]):

    def __init__(self, source: 'DBSource', table_name: str, model_type: PEEWEE_MODEL_TYPE, indexes=None):
        self._source = source
        self._table_name = table_name
        self._indexes = indexes
        self._origin_model_type = model_type
        self._stat = _PeeweeModelState()


    def _get_state_model(self):
        if self._stat.model is None:
            table_name = self._table_name
            indexes_desc = self._indexes
            class _Inner_Model(self._origin_model_type):
                class Meta:
                    database = None
                    db_table = table_name
                    indexes = indexes_desc
            self._stat.model =_Inner_Model
        return self._stat.model
        # peewee_model = self._stat.model
        # if peewee_model is None:
        #     ##新线程，则使用新的Model
        #     class SubModel(self._origin_model_type):
        #         pass
        #     self._stat.model = SubModel
        #     peewee_model = self._stat.model
        # assert peewee_model is not None
        # return peewee_model

    def use_peewee(self,read_only = False) -> typing.ContextManager[PEEWEE_MODEL_TYPE]:
        return PeeweeModelWithStatement(self,read_only)


class DBSource:


    def __init__(self,**setting):
        self._peewee_db_pool:_PeeweeConnectPool = None
        self._redis_db_pool = None
        peewee_setting = setting.get("peewee")
        self._lock = threading.RLock()
        self._peewee_sources:typing.Dict[str,PeeweeModelSource] = {}
        if peewee_setting:
            type = peewee_setting['type']
            connect_param = peewee_setting['connect_param']
            pool_param = peewee_setting['pool_param']
            if not connect_param:
                connect_param = {}
            if type == 'mysql':
                database = peewee_setting['database']
                connect_builder = _MysqlConnectBuilder(database,**connect_param)
            elif type== 'sqlite':
                file_path = peewee_setting['file_path']
                connect_builder = _SqliteConnectBuilder(file_path,**connect_param)
            else:
                raise RuntimeError(f"unkonw peewee database type: {type}")

            if not pool_param:
                pool_param = {}
            self._peewee_db_pool = _PeeweeConnectPool(connect_builder, **pool_param)

    def name(self):
        raise NotImplementedError()

    def getDescription(self,use_markdown = True):
        """
        返回数据库源的描述。
        :param use_markdown: 是否使用MarkDown方式。
        :return:
        """
        desc = ""

        if use_markdown:
            if self._peewee_db_pool:
                desc += f"+ peewee数据源:{self._peewee_db_pool.builder.get_db_path_desc()}"
            else:
                desc += f"+ peewee数据源: 没有配置"
            if self._redis_db_pool:
                desc += f"\n+ redis数据源:{self._redis_db_pool.builder.get_db_path_desc()}"
            else:
                desc += f"\n+ redis数据源: 没有配置"
        else:
            if self._peewee_db_pool:
                desc += f"peewee数据源:{self._peewee_db_pool.builder.get_db_path_desc()}"
            else:
                desc += f"peewee数据源: 没有配置"
            if self._redis_db_pool:
                desc += f"\nredis数据源:{self._redis_db_pool.builder.get_db_path_desc()}"
            else:
                desc += f"\nredis数据源: 没有配置"
        return desc

    def init_peewee(self, table_name: str, model_type: PEEWEE_MODEL_TYPE, indexes=None, ignore_inited=True)->PeeweeModelSource[PEEWEE_MODEL_TYPE]:
        if not self._peewee_db_pool:
            raise RuntimeError("not config peewee!!")
        with self._lock:
            model_source = self._peewee_sources.get(table_name)
            if model_source is not None:
                if ignore_inited:
                    return model_source
                raise RuntimeError(f"table_name:{table_name} has inited!!")

            model_source = PeeweeModelSource(self,table_name,model_type,indexes)
            self._peewee_sources[table_name] = model_source
            with model_source.use_peewee() as db_model:
                db_model.create_table()

        return model_source

    def use_connect(self,read_only = False)->PeeweeDatabaseWithStatementSingle:
        """
        配合with使用，获得基于peewee的数据库连接对象peewee.Database
        :return:
        """
        if not self._peewee_db_pool:
            raise RuntimeError("not config peewee!!")
        session = self._peewee_db_pool.get_connect_session(read_only)
        return session

    def peewee_source(self, table_name: str, model_type: PEEWEE_MODEL_TYPE = Type[peewee.Model]) -> PeeweeModelSource[
        PEEWEE_MODEL_TYPE]:
        if not self._peewee_db_pool:
            raise RuntimeError("not config peewee!!")
        model_source = self._peewee_sources.get(table_name)
        if not model_source:
            raise RuntimeError(f"must init_peewee() fisrt, table_name: {table_name}")
        return model_source

    def use_peewee(self, table_name: str, model_type: PEEWEE_MODEL_TYPE = Type[peewee.Model], read_only=False) -> \
    typing.ContextManager[PEEWEE_MODEL_TYPE]:
        """
        :return:
        """
        model_source = self.peewee_source(table_name, model_type)
        return model_source.use_peewee(read_only)


    def use_redis(self,read_only = False):
        """
        配合with使用，获得Redis数据库连接对象
        :return:
        """
        pass





    # def _dump(self):
    #     if self.peewee_db_pool:
    #         print(f"use size:{len()}")

class DBSourceBuilder:

    def __init__(self):
        self.settings = {}

    def redis(self):
        return self

    def peewee_use_sqlite(self,file_path,connect_param:{}=None,pool_param:{} = None):
        self.settings['peewee'] = {
            "type": "sqlite",
            "file_path": file_path,
            "connect_param": connect_param,
            "pool_param": pool_param
        }
        return self

    def peewee_use_mysql(self,database:str,connect_param:{},pool_param:{} = None):
        """
        使用mysql作为peeweee的数据库连接

        :param database: 数据库名称
        :param connect_param: 相关msyql的连接对象。

        :param pool_param: 连接池的设置。
            time_out：数据连接池满了之后，等待连接超时秒数。默认不等待
            live_time:每次连接的生存时间，时间过期之后的连接将被剔除连接池并关闭。默认是：60*60*2秒，即2小时
            max_connection: 链接池的最大连接数，包含两种模式：
                1) max_connection = 20，标示read=False的最大连接数为20，read=True的最大连接数为0
                2) max_connection = [1,20]，标示read=False的最大连接数为20，read=True的最大连接数为20
           {
            "time_out": 5.5,
             "live_time": 60 * 60 * 2,
              "max_connection": [1, 20]
          }

        :return:
        """
        self.settings['peewee'] = {
            "type":"mysql",
            "database":database,
            "connect_param":connect_param,
            "pool_param": pool_param
        }
        return self



    def build(self)->DBSource:
        return DBSource(**frozendict(self.settings))



class _PeeweeConnectSession(typing.ContextManager[peewee.Database]):

    def __init__(self,id:int,read_only:bool,pool, conn: peewee.Database):
        self.id = id
        self.read_only = read_only
        self._conn = conn
        self._pool = pool
        self._session_lock = threading.RLock()
        self._has_enter = False
        self.create_timestamp = time.time()  ##创建的开始时间
        self.checkout_timestamp = 0  ##上次使用的时间。
        self.occur_error = False
        self.life_end = False

    def __enter__(self):
        if self._session_lock:
            if self._has_enter:
                raise RuntimeError("Multiple with nested error")  ## 不支持多重内嵌with
            self._has_enter = True
        self.checkout_timestamp = time.time()
        return self._conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        with self._session_lock:
            if not self._has_enter:
                raise RuntimeError("shoul't not run here")  ## 不支持多重内嵌with
            self._has_enter = False

        if exc_type is not None:
            self.occur_error = True
        self._pool.release_connect_session(self)  ## 释放session

    def destroy(self):
        self._conn = None
        self._pool = None

class _PeeweeConnectBuilder:

    def create_connect(self,id,read_only)->peewee.Database:
        raise NotImplementedError()

    def destroy_connect(self,id,read_only:bool,is_error:bool,is_life_end:bool):
        pass

    def get_db_path_desc(self):
        raise NotImplementedError()


class _PeeweeConnectPool:

    def __init__(self,builder:_PeeweeConnectBuilder,**kwargs):
        self.builder = builder
        self.time_out = kwargs.get('time_out')
        self.max_live_time = kwargs.get('live_time')
        max_connection = kwargs.get('max_connection')
        if not self.time_out:
           self.time_out = 5
        if not self.max_live_time:
            self.max_live_time = 60 * 60 * 2  # 即2小时
        self.max_live_time = self.max_live_time + random.random()* 10   ## 加个0-1随机数，避免扎堆创建销毁
        if max_connection:
            if isinstance(max_connection,list):
                self.max_connection_rand_and_write = max_connection[0]
                try:
                    self.max_connection_readonly = max_connection[1]
                except:
                    self.max_connection_readonly = 0
            else:
                self.max_connection_readonly = 0
                self.max_connection_rand_and_write = max_connection
        else:
            self.max_connection_readonly = 0
            self.max_connection_rand_and_write = 20
        self._lock = threading.RLock()
        self.id_token = 0

        self._sessions_readonly = []
        self._in_use_readonly = {}
        self._sessions_read_write = []
        self._in_use_read_write = {}


    def get_connect_session(self,read_only:bool)->_PeeweeConnectSession:
        closing_session = []
        _new_session: _PeeweeConnectSession = None  ## 新建的session
        _start_time = time.time()
        session = None
        wait_time_second = 0
        while True:
            no_use_session = None
            with self._lock:
                session_list = self._sessions_readonly if read_only else self._sessions_read_write
                in_use = self._in_use_readonly if read_only else self._in_use_read_write
                try:
                    __no_user__, no_use_session = heapq.heappop(session_list)
                except IndexError:
                    ### 看是否需要创建。
                    in_use_size = len(in_use)
                    max_connect_size = self.max_connection_readonly if read_only else self.max_connection_rand_and_write
                    if in_use_size < max_connect_size:
                        ## 线程池还有容量，创建一个。
                        self.id_token += 1
                        new_id = self.id_token
                        _new_conn = self.builder.create_connect(new_id,read_only)
                        _new_session = _PeeweeConnectSession(new_id, read_only, self, _new_conn)
                        in_use[_new_session.id] = _new_session
                        session = _new_session
                        if _DEBUG:
                            logger.debug('No Session available in pool.')
                        break
                    ### 等待释放connection
                if no_use_session:
                    if  self._can_reuse(no_use_session):
                        session = no_use_session
                        break
                    in_use.pop(no_use_session.id)
                    closing_session.append(no_use_session)
            session = None
            ## 等待资源释放
            if self.time_out < 0.01:
                raise RuntimeError('Exceeded maximum connections.')
            now = time.time()
            total_wait_time = now - _start_time
            if total_wait_time > self.time_out:
                raise RuntimeError('Exceeded maximum connections and connect time out.')
            wait_time_second += 0.1
            if wait_time_second >= 1:
                wait_time_second = 1
            time.sleep(wait_time_second)

        assert session is not None
        ### 允许新建一个session,新建session
        if _new_session:
            try:
                _new_session._conn.connect()
                session = _new_session
                if _DEBUG:
                    logger.debug(f'Created new session {_new_session.id}.')
            except Exception:
                with self._lock:
                    in_use.pop(session.id)
                closing_session.append(_new_session)
                raise RuntimeError("connect db error")
        # else:
        #     if _DEBUG:
        #         logger.debug(f'Use pool session {session.id}.')

        ##清理不要的的Session
        if closing_session:
            for _no_use_session in closing_session:
                self._destroy_session(_no_use_session)
        return session

    def release_connect_session(self,session:_PeeweeConnectSession):
        _destroy_session = False
        with self._lock:
            read_only = session.read_only
            if self._can_reuse(session):
                session_list = self._sessions_readonly if read_only else self._sessions_read_write
                ## 继续放在线程池里面。
                heapq.heappush(session_list, (session.checkout_timestamp, session))
                if _DEBUG:
                    logger.debug(f'Session {session.id} push back. pool_size: {len(session_list)}')

            else:
                in_use = self._in_use_readonly if read_only else self._in_use_read_write
                in_use.pop(session.id)
                _destroy_session = True
        if _destroy_session:
            self._destroy_session(session)


    # def _create_connect_and_wait(self,read_only:bool)->peewee.Database:
    #     need_wait = self.time_out > 0
    #     _start_time = None
    #     if need_wait:
    #         _start_time = time.time()
    #     while True:
    #         in_use = self._in_use_readonly if read_only else self._in_use_read_write
    #         max_connect_size = self.max_connection_readonly if read_only else self.max_connection_rand_and_write
    #         in_use_size = len(in_use)
    #         if in_use_size >= max_connect_size:
    #             ### 是否继续等待
    #             if not need_wait:
    #                 raise RuntimeError('Exceeded maximum connections.')
    #             now = time.time()
    #             total_wait_time = now - _start_time
    #             if total_wait_time < self.time_out:
    #                 time.sleep(0.1)
    #             else:
    #                 raise RuntimeError('Exceeded maximum connections and connect time out.')
    #         else:
    #             return self.builder.create_connect(read_only)

    def _can_reuse(self,session:_PeeweeConnectSession):
        live_time = time.time() - session.create_timestamp
        if session.occur_error:
            if _DEBUG:
                logger.debug(f'Session {session.id} was closed because of error occur.')
            return False
        if live_time >= self.max_live_time:
            # 到底生命周期，关闭连接
            session.life_end = True
            if _DEBUG:
                logger.debug(f'Session {session.id} was closed because of life end.')
            return False
        return True

    def _destroy_session(self, session: _PeeweeConnectSession):
        try:
            session._conn.close()
        except Exception:
            pass
        session.destroy()
        self.builder.destroy_connect(session.id,session.read_only,session.occur_error,session.life_end)


class _SqliteConnectBuilder(_PeeweeConnectBuilder):

    def __init__(self,file_path,**kwargs):
        self.path = file_path
        self.setting = dict(kwargs)
        self.setting['check_same_thread'] = False

    def create_connect(self,id,read_only)->peewee.Database:
        filePath = self.path
        if read_only:
            filePath = f"file:{filePath}?mode=ro"
            return peewee.SqliteDatabase(filePath, uri=True, check_same_thread=False,thread_safe=False,**self.setting)
        else:
            return peewee.SqliteDatabase(filePath, thread_safe=False,**self.setting)

    def get_db_path_desc(self):
        return f"sqlite:{self.path}"

class _MysqlConnectBuilder(_PeeweeConnectBuilder):

    def __init__(self,database,**kwargs):
        self.database = database
        self.setting = dict(kwargs)

    def create_connect(self,id,read_only)->peewee.Database:
        report.log_info(f"==>创建连接,id:{id}, read_only: {read_only}",tag="MYSQL")
        return peewee.MySQLDatabase(self.database, thread_safe=False, **self.setting)

    def destroy_connect(self,id,read_only:bool,is_error:bool,is_life_end:bool):
        report.log_info(f"<==销毁连接,id:{id}, read_only: {read_only}，is_error:{is_error},is_life_end:{is_life_end}",tag="MYSQL")

    def get_db_path_desc(self):
        host = self.setting.get("host")
        port = self.setting.get("port")
        user= self.setting.get("user")
        return f"mysql:/{host}:{port}?user={user}"
